#!/usr/bin/env python
# encoding: utf-8
# Author: 潘世友
# Email: stpann@163.com

__note__ = """
相关性计量：资产相关性或违约相关性。
银行贷款资产组合的风险与经济资本计量可以看作是从单纯独立地考察单笔贷款风险上升到全盘考虑整个资产组合风险的实施过程，\
在这个过程中，一个最为重要的参数就是“违约相关性”，单个贷款风险和资本需求不能直接简单加总得到整个资产组合风险和资本 \
需求的本质原因就是因为违约相关性的存在。一个贷款的违约以致发生损失（或评级迁移）不意味着另一个贷款同时违约（或评级 \
迁移）发生损失，即违约相关性或称损失的相关性不为1；但是当一个贷款违约发生损失（或评级迁移）时，另一个贷款的违约概 \
率也可能会增加，即两笔贷款违约（或评级迁移）和发生损失是不独立的，相关性不为0。
违约相关性的估计有两种方法：按照违约相关性的定义，最直接的方法是，从银行内部数据出发直接估计客户两两联合违约概率从 \
而得出两两违约相关性。然而，实践发现该方法的使用存在一个普遍的难题，即需要多年大量违约数据，而总所周知银行的违约数 \
据是不可能快速积累的。
随着内评法在银行业的推广，越来越多国内银行认识到违约数据的重要性，银行已经逐步建立数据的有效收集机制，但考虑到银行 \
的实际业务开展情况，两类实践导致违约相关数据依然与理想状况存在较大的误差。
第一类实践是几乎所有银行的实际业务政策和业务发展中都会通过有效的风险管理、控制客户质量以保持全行平均违约概率或不良 \
率在一定范围内，以确保银行的盈利能力与持续经营。从这个角度说，越是好银行其违约数据可能就越少。第二类实不同行业客户 \
数量的天然差异导致部分行业的客户数据严重不足，尤其是违约数据，例如金融业的违约客户数量基本为0，因此从这个角度来说，\
部分行业的违约相关性天然的无法利用历史数据进行合理准确的估计。]，因此尝试通过外部金融市场上较为充足的数据估计违约 \
相关性的方法应运而生。即通过客户资产相关性的估计，然后基于假设“客户资产收益率下跌至某一临界值时触发违约”间接获取客 \
户的违约相关性。
在盯市模式下，违约相关性估计将要扩展为评级迁移相关性，评级迁移相关性是指客户同时迁移到某种级别的相关性，其中迁移到 \
违约级别即为违约相关性。此时如果使用内部数据估计评级迁移相关性，那么除了要积累违约数据外还需要积累大量的评级迁移数 \
据；如果使用资产相关性推导评级迁移相关性的话，则需要进一步假设“客户资产收益率下跌（或上升）至一系列临界值时就会触 \
发客户评级迁移到临界值对应的评级”。
"""

import warnings
from math import sqrt
from typing import Literal

import numpy as np
import pandas as pd
from scipy.optimize import root
from scipy.stats import norm, multivariate_normal

from rmwheels.ec import rwa
from rmwheels.utils.check import check_para, check_data


def _calc_rho_d(pdv1: float, pdv2: float, rho_a: float) -> float:
    x = np.array([norm.ppf(pdv1), norm.ppf(pdv2)])

    cov = np.array([[1, rho_a], [rho_a, 1]])
    jpd = multivariate_normal.cdf(x=x, mean=np.array([0, 0]), cov=cov)
    rho_d = (jpd - pdv1 * pdv2) / sqrt(pdv1 * (1 - pdv1) * pdv2 * (1 - pdv2))
    return float(rho_d)


def _calc_rho_a(pdv1: float, pdv2: float, rho_d: float) -> float:
    # 特殊情况处理，当违约概率[pd]为0或1、违约相关性为0或1时，直接返回相关性
    if rho_d == 0 or pdv1 == 0 or pdv2 == 0:
        return 0
    if pdv1 == 1 or pdv2 == 1:
        return 1
    if rho_d in [-1, 1]:
        return rho_d

    def opt_func(args):
        """ 根据违约相关性与资产相关性之间的理论关系反算资产相关性 """
        if isinstance(args, np.ndarray):
            args = args.item()
        args = float(args)

        cov = np.array([[1, args], [args, 1]], dtype=np.float64)
        x = np.array([norm.ppf(pdv1), norm.ppf(pdv2)], dtype=np.float64)
        mean = np.array([0, 0], dtype=np.float64)

        try:
            # 尝试使用更新的接口
            jpd, _ = multivariate_normal.cdf(x=x, mean=mean, cov=cov, maxpts=100000)
        except TypeError:
            # 回退到旧的接口
            jpd = multivariate_normal.cdf(x=x, mean=mean, cov=cov)

        # 确保计算过程中的类型一致性
        numerator = jpd - pdv1 * pdv2
        denominator = sqrt(pdv1 * (1 - pdv1) * pdv2 * (1 - pdv2))
        rho_d2 = numerator / denominator

        return rho_d2 - rho_d

    try:
        rho_a = root(opt_func, x0=0.2, method='hybr')
        return float(rho_a.x[0])
    except ValueError:
        raise Exception('给定的违约相关性下，无法正确计算资产相关性，请确保违约率和违约相关性合理')


@check_para({'pdv1': (0, 1),
             'pdv2': (0, 1),
             'exp1': ['cor', 'fin', 'sme', 'mor', 'rec', 'rtl'],
             'exp2': ['cor', 'fin', 'sme', 'mor', 'rec', 'rtl'],
             'rho_a': (-1, 1),
             'rho_d': (-1, 1),
             })
def convert_df_as(pdv1: float,
                  pdv2: float,
                  exp1: Literal['cor', 'fin', 'sme', 'mor', 'rec', 'rtl'] = 'cor',
                  exp2: Literal['cor', 'fin', 'sme', 'mor', 'rec', 'rtl'] = 'cor',
                  rho_a: float | None = None,
                  rho_d: float | None = None) -> tuple[float | None, float | None]:
    """
    根据输入的违约率数值、违约相关性、资产相关性数据转换违约相关性、资产相关性结果
    存在两种基本情况，已知资产相关性或违约相关性，或未知两债项之间的相关性。
    * 已知资产相关性或违约相关性时，利用公式计算违约相关性、资产相关性
    * 对于债项之间的相关性无法确定的时候，利用BASEL公式计算违约相关性、资产相关性


    Parameters
    ----------
    pdv1: float，债项A违约率
    pdv2: float，债项B违约率
    exp1: Literal['cor', 'fin', 'sme', 'mor', 'rec', 'rtl'] = 'cor'，债项A的敞口类型，默认值为'cor'
        - 'cor': 主权、一般公司风险暴露
        - 'fin': 金融机构风险暴露
        - 'sme': 中小企业风险暴露
        - 'mor': 个人住房抵押贷款
        - 'rec': 合格循环零售贷款
        - 'rtl': 零售风险暴露
    exp2: 同exp1
    rho_a: float，资产相关性数值，取值范围为(-1,1)
    rho_d: float，违约相关性数值，取值范围为(-1,1)

    Return
    -------
    rho_a: float，资产相关性数值
    rho_d: float，违约相关性数值

    Examples
    --------
    >>> # 未知两客户之间的相关性，利用BASEL公式计算违约相关性、资产相关性
    >>> ra, rd = convert_df_as(0.1, 0.02)
    >>> print(ra, rd)
    0.14081968761859825 0.03395179012079394

    >>> # 已知两客户之间的资产相关性，计算违约相关性
    >>> ra, rd = convert_df_as(0.1, 0.02, rho_a = 0.23)
    >>> print(ra, rd)
    0.23 0.061381274623794564

    >>> # 已知两客户之间的违约相关性，计算资产相关性
    >>> ra, rd = convert_df_as(0.1, 0.02, rho_d = 0.06)
    >>> print(ra, rd)
    0.22585442117941054 0.06
    """

    if rho_a is not None and rho_d is not None:
        warnings.warn("资产相关性、违约相关性已经给定，无需计算，直接返回")
        return rho_a, rho_d

    if pdv1 == 0 or pdv2 == 0 or pdv1 == 1 or pdv2 == 1:
        warnings.warn("违约率为0或1时，相关性无法计算，直接返回1")
        return 1, 1

    if rho_a is not None:
        return rho_a, _calc_rho_d(pdv1, pdv2, rho_a)
    if rho_d is not None:
        return _calc_rho_a(pdv1, pdv2, rho_d), rho_d

    if rho_a is None and rho_d is None:
        rho1 = rwa.rwa_r(pdv1, exp1)
        rho2 = rwa.rwa_r(pdv2, exp2)
        rho_a = sqrt(rho1 * rho2)
        return rho_a, _calc_rho_d(pdv1, pdv2, rho_a)


def convert_dfm_asm(list_pdv: list[float] | np.ndarray,
                    list_exposure: list[Literal['cor', 'fin', 'sme', 'mor', 'rec', 'rtl']] | None = None,
                    matr_rho_a: np.ndarray | None = None,
                    matr_rho_d: np.ndarray | None = None) -> tuple[np.ndarray, np.ndarray]:
    """
    根据输入的违约率数组、违约相关性矩阵、资产相关性矩阵转换违约相关性矩阵、资产相关性矩阵

    Parameters
    ----------
    list_pdv: list[float]，债项的违约率清单
    list_exposure: list[Literal['cor', 'fin', 'sme', 'mor', 'rec', 'rtl']] | None = None
        敞口类型清单，空或者未指定pdv的情况时，默认为'cor'
        - 'cor': 主权、一般公司风险暴露
        - 'fin': 金融机构风险暴露
        - 'sme': 中小企业风险暴露
        - 'mor': 个人住房抵押贷款
        - 'rec': 合格循环零售贷款
        - 'rtl': 零售风险暴露
    matr_rho_a: np.ndarray | None = None，资产相关性矩阵，取值范围为(-1,1)
    matr_rho_d: np.ndarray | None = None，违约相关性矩阵，取值范围为(-1,1)

    Return
    -------
    rho_a_matr: np.ndarray，资产相关性矩阵
    rho_d_matr: np.ndarray，违约相关性矩阵

    Examples
    --------
    >>> # 未知两客户之间的相关性，利用BASEL公式计算违约相关性、资产相关性
    >>> ra, rd = convert_dfm_asm([0.1, 0.02])
    >>> print(ra)
    array([[0.12080855, 0.14081969],
           [0.14081969, 0.16414553]])
    >>> print(rd)
    array([[0.04549638, 0.03395179],
           [0.03395179, 0.0273688 ]])

    >>> # 已知两客户之间的资产相关性，计算违约相关性
    >>> ra, rd = convert_dfm_asm([0.1, 0.02], matr_rho_a=np.array([[0.16, 0.23], [0.23, 0.12]]))
    >>> print(rd)
    array([[0.06207919, 0.06138127],
           [0.06138127, 0.01834915]])

    >>> # 已知两客户之间的违约相关性，计算资产相关性
    >>> ra, rd = convert_dfm_asm([0.1, 0.02], matr_rho_d=np.array([[0.04, 0.06], [0.06, 0.01]]))
    >>> print(ra)
    array([[0.10732211, 0.22585442],
           [0.22585442, 0.07201665]])
    """
    # 检查list_pdv的类型是否符合要求
    if not isinstance(list_pdv, (list, np.ndarray)):
        raise TypeError("list_pdv类型错误，必须是list或np.ndarray类型")
    if isinstance(list_pdv, list) and len(list_pdv) <= 1:
        raise ValueError("list_pdv长度必须大于1")
    if isinstance(list_pdv, np.ndarray):
        if np.ndim(list_pdv) != 1 or len(list_pdv) <= 1:
            raise ValueError("输入的list_pdv必须是一维数组且长度大于1")

    # 检查matr_rho_a、matr_rho_d是否为对称矩阵
    if matr_rho_a is not None and not np.allclose(matr_rho_a, matr_rho_a.T):
        raise ValueError("matr_rho_a必须是对称矩阵")
    if matr_rho_d is not None and not np.allclose(matr_rho_d, matr_rho_d.T):
        raise ValueError("matr_rho_d必须是对称矩阵")

    # 检查list_exposure的长度是否与list_pdv一致
    if list_exposure is None:
        list_exposure = ['cor'] * len(list_pdv)
    elif len(list_exposure) < len(list_pdv):
        list_exposure.extend(['cor'] * (len(list_exposure) - len(list_pdv)))

    n_dim = len(list_pdv)

    def check_input_mat(input_mat):
        if input_mat is None:
            input_mat = np.zeros((n_dim, n_dim))
            input_mat[:] = None
        else:
            if input_mat.shape[0] < n_dim:
                r_dim, c_dim = n_dim - input_mat.shape[0], input_mat.shape[1]
                patch = np.zeros((r_dim, c_dim))
                patch[:] = None
                input_mat = np.r_[input_mat, patch]
            if input_mat.shape[1] < n_dim:
                r_dim, c_dim = input_mat.shape[0], n_dim - input_mat.shape[1]
                patch = np.zeros((r_dim, c_dim))
                patch[:] = None
                input_mat = np.c_[input_mat, patch]
        return input_mat

    matr_rho_a = check_input_mat(matr_rho_a)
    matr_rho_d = check_input_mat(matr_rho_d)

    matr_rho_a_out: np.ndarray = np.zeros((n_dim, n_dim))
    matr_rho_d_out: np.ndarray = np.zeros((n_dim, n_dim))
    for i in range(n_dim):
        for j in range(n_dim):
            if np.isnan(matr_rho_a[i, j]):
                rho_a = None
            else:
                rho_a = matr_rho_a[i, j]
            if np.isnan(matr_rho_d[i, j]):
                rho_d = None
            else:
                rho_d = matr_rho_d[i, j]
            matr_rho_a_out[i, j], matr_rho_d_out[i, j] = convert_df_as(pdv1=list_pdv[i],
                                                                       pdv2=list_pdv[j],
                                                                       exp1=list_exposure[i],
                                                                       exp2=list_exposure[j],
                                                                       rho_a=rho_a,
                                                                       rho_d=rho_d)
    return matr_rho_a_out, matr_rho_d_out


class InternalCorrMatrix:
    """
    根据内部违约数据计算违约相关性矩阵、或通过莫顿公式转换为资产相关性矩阵

    Parameters
    ----------
    data：pandas.DataFrame，必须保证有字段“分组”、“日期”、“客户数”、“违约客户数”
        - group：分组名称
        - date: 日期字段
        - customer: 客户数字段
        - default: 违约客户数字段

    >>> from rmwheels.datasets import load_data
    >>> data = load_data("default_history")
    >>> icm = InternalCorrMatrix(data, group='分组', date='年份', customers='客户数', defaults='违约数')
    >>> print(icm.data.head(5))
        年份   分组   客户数  违约数
    0  2018  a01   600   31
    1  2017  a01   701   35
    2  2016  a01   994   61
    3  2015  a01   631   43
    4  2014  a01  1271   34
    """

    def __init__(self,
                 data: pd.DataFrame,
                 weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'simple'] = 'pair',
                 sample_thresh: int = 20,
                 neg_rho: float | None = None,
                 agg: Literal['mean', 'sample'] = 'mean',
                 **kwargs: dict[str, str]):
        self.data = data
        self.sample_thresh = sample_thresh
        self.neg_rho = neg_rho
        self.agg = agg

        # check the input data
        columns: list = ['group', 'date', 'customers', 'defaults']
        data_: pd.DataFrame = check_data(data, columns, **kwargs)

        # initiate the source matrix
        data_pvt = pd.pivot_table(data_, values=['customers', 'defaults'], index='group', columns='date', aggfunc='sum')
        self.__customers_array__: np.ndarray = np.array(data_pvt['customers'])
        self.__defaults_array__: np.ndarray = np.array(data_pvt['defaults'])
        self.__group__: list = list(data_pvt.index)

    @staticmethod
    def _weight(arr: np.ndarray) -> np.ndarray:
        return arr / arr.sum()

    def _weight_simple(self, arr_n1: np.ndarray, arr_n2: np.ndarray) -> np.ndarray:
        """ 计算简单的权重 """
        arr_ = np.array([1] * self.__ndim__)
        arr_w1 = self._weight(arr_)
        arr_w2, arr_w3 = arr_w1, arr_w1
        return arr_w1, arr_w2, arr_w3

    def _weight_pair(self, arr_n1: np.ndarray, arr_n2: np.ndarray) -> np.ndarray:
        """ 计算以客户数的权重 """
        arr_w1 = self._weight(arr_n1 * (arr_n1 - 1))
        arr_w2 = self._weight(arr_n2 * (arr_n2 - 1))
        arr_w3 = self._weight(arr_n1 * arr_n2)
        return arr_w1, arr_w2, arr_w3

    def _weight_time(self, arr_n1: np.ndarray, arr_n2: np.ndarray, time_lag: int = 1) -> np.ndarray:
        """ 计算以时间的权重 """
        arr_ = np.arange(1, self.__ndim__ + 1) + time_lag
        arr_ = self._weight(1 / np.log(arr_))
        arr_w1, arr_w2, arr_w3 = arr_, arr_, arr_
        return arr_w1, arr_w2, arr_w3

    def _weight_time_pair(self, arr_n1: np.ndarray, arr_n2: np.ndarray, time_lag: int = 1) -> np.ndarray:
        """ 计算以时间的权重 """
        arr_ = np.arange(1, self.__ndim__ + 1) + time_lag
        arr_ = 1 / np.log(arr_)

        arr_w1 = self._weight(arr_n1 * (arr_n1 - 1) * arr_)
        arr_w2 = self._weight(arr_n2 * (arr_n2 - 1) * arr_)
        arr_w3 = self._weight(arr_n1 * arr_n2 * arr_)

        return arr_w1, arr_w2, arr_w3

    def _weight_time_ln_pair(self, arr_n1: np.ndarray, arr_n2: np.ndarray, time_lag: int = 1) -> np.ndarray:
        arr_ = np.arange(1, self.__ndim__ + 1) + time_lag
        arr_ = 1 / np.log(arr_)

        arr_w1 = self._weight(np.log(arr_n1 * (arr_n1 - 1)) * arr_)
        arr_w2 = self._weight(np.log(arr_n2 * (arr_n2 - 1)) * arr_)
        arr_w3 = self._weight(np.log(arr_n1 * arr_n2) * arr_)

        return arr_w1, arr_w2, arr_w3

    def _weight_array(self,
                      weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'simple'],
                      arr_n1: np.ndarray,
                      arr_n2: np.ndarray,
                      time_lag: int = 1) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        # 确定计量相关性的权重
        if weight == "pair":
            return self._weight_pair(arr_n1, arr_n2)
        if weight == "time":
            return self._weight_time(arr_n1, arr_n2, time_lag)
        if weight == 'time_pair':
            return self._weight_time_pair(arr_n1, arr_n2, time_lag)
        if weight == 'time_ln_pair':
            return self._weight_time_ln_pair(arr_n1, arr_n2, time_lag)
        if weight == "simple":
            return self._weight_simple(arr_n1, arr_n2)

        raise NotImplementedError(f"weight参数错误，{weight}不存在")

    def _inner_rho(self,
                   weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'simple'],
                   arr_n1: np.ndarray,
                   arr_p1: np.ndarray,
                   sample_thresh: int,
                   neg_rho: float | None = None,
                   time_lag: int = 1) -> tuple[float, float]:
        """ 计算内部违约率数组 """
        arr_w1, *_ = self._weight_array(weight, arr_n1, arr_n1, time_lag)

        mask = (arr_n1 >= thresh)
        arr_n1, arr_p1, arr_w1 = arr_n1[mask], arr_p1[mask], self._weight(arr_w1[mask])

        arr_pd1 = arr_w1 * (arr_p1 / arr_n1)
        pd_k = arr_pd1.sum()
        arr_jpd = arr_w1 * (arr_p1 / arr_n1) * ((arr_p1 - 1) / (arr_n1 - 1))
        arr_jpd[np.where(arr_p1 == 0)] = 0  # 修正违约数为0的情况
        jpd_k = arr_jpd.sum()

        # 当pd_k为0时，rho为0
        if pd_k * (1 - pd_k) == 0:
            return 0, 0

        rho = (jpd_k - pd_k * pd_k) / (pd_k * (1 - pd_k))
        # 当rho为NaN时，返回0
        if np.isnan(rho):
            return 0, 0
        # 当rho为负数时，neg_rho作为违约相关性
        if neg_rho is not None and rho < 0:
            return convert_df_as(pd_k, pd_k, rho_d=neg_rho)
        # 当rho为正数时，计算资产相关性
        return convert_df_as(pd_k, pd_k, rho_d=rho)

    def _outer_rho(self,
                   weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'simple'],
                   arr_n1: np.ndarray,
                   arr_n2: np.ndarray,
                   arr_p1: np.ndarray,
                   arr_p2: np.ndarray,
                   sample_thresh: int,
                   neg_rho: float | None = None,
                   time_lag: int = 1) -> tuple[float, float]:
        """ 计算违约率数组 """
        arr_w1, arr_w2, arr_w3 = self._weight_array(weight, arr_n1, arr_n2, time_lag)

        mask = (arr_n2 >= thresh) * (arr_n1 >= thresh)
        arr_n1, arr_n2 = arr_n1[mask], arr_n2[mask]
        arr_p1, arr_p2 = arr_p1[mask], arr_p2[mask]
        arr_w1 = self._weight(arr_w1[mask])
        arr_w2 = self._weight(arr_w2[mask])
        arr_w3 = self._weight(arr_w3[mask])

        arr_pd1 = arr_w1 * (arr_p1 / arr_n1)
        pd_k = arr_pd1.sum()
        arr_pd2 = arr_w2 * (arr_p2 / arr_n2)
        pd_l = arr_pd2.sum()
        arr_jpd = arr_w3 * (arr_p1 / arr_n1) * (arr_p2 / arr_n2)
        jpd_kl = arr_jpd.sum()

        # 当pd_k或pd_l为0时，rho为0
        if pd_k * (1 - pd_k) * pd_l * (1 - pd_l) == 0:
            return 0, 0
        # 计算相关性
        rho = (jpd_kl - pd_k * pd_l) / sqrt(pd_k * (1 - pd_k) * pd_l * (1 - pd_l))
        # 当rho为NaN时，返回0
        if np.isnan(rho):
            return 0, 0
        # 当rho为负数时，neg_rho作为违约相关性
        if neg_rho is not None and rho < 0:
            return convert_df_as(pd_k, pd_l, rho_d=neg_rho)
        # 当rho为正数时，计算资产相关性
        return convert_df_as(pd_k, pd_l, rho_d=rho)

    def default_corr_matr(self,
                          weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'simple'] = self.weight,
                          sample_thresh: int = self.sample_thresh,
                          neg_rho: float | None = self.neg_rho,
                          agg: Literal['mean', 'sample'] = self.agg):
        """
        给定违约统计数据后计算其对应的违约相关性矩阵以及相关性总体的平均值

        Parameters
        ----------
        weight: str，历史数据权重参数，{'pair','time','time_pair','time_ln_pair','average'}，默认为'pair'
            - pair: 使用样本的配对数作为权重
            - time: 根据时间长短确定权重，默认[3,3,3,2,2,2,1,...]，7年前的权重统一设置为1，若希望改变该权重列表，使用
                    time_weight进行自定义，其顺序是按历史时间由近至远进行分配
            - time_pair: 综合样本配对数与时间长短确定权重
            - time_ln_pair: 综合样本配对数的ln处理与时间长短确定权重
            - average: 每一年的相关性进行简单平均
        sample_thresh: int，计算违约相关性的历史数据最小样本量的设定，默认为30
        neg_rho: float，违约相关性计算结果为负值的时候进行人为修改的取值，默认不修改
        time_weight: list，当权重参数为time、time_pair、time_ln_pair的时候，确定的时间权重列表，默认为[3,3,3,2,2,2,1]
        aggfunc: str，总体相关性平均值的计量方法，默认值为mean
            - mean: 基于违约相关性矩阵，组间相关性为所有与目标分组之间的相关性的平均值
            - sample: 组间相关性是基于将所有非目标分组汇总结果与目标分组，根据相关性计算公式重新计算

        Return
        ------
        dsn_rho: pandas.DataFrame, 包含组名的相关性矩阵数据
        dsn_rho_avg: pandas.DataFrame, 包含组名的相关性平均结果数据

        Examples
        --------
        >>> from rmwheels.datasets import load_data
        >>> data = load_data("default_history")
        >>> x1 = CorrInternalData(data, group='分组', date='年份', cust='客户数', default='违约数')
        >>> # 利用默认设置计算违约相关性、违约相关性平均值
        >>> dfm11, dfm12 = x1.rho_default_matrix()
        >>> dfm11
                  a01       a02       a03  ...       a10       a11       a12
        a01  0.015710  0.014449  0.015605  ...  0.021983  0.016301  0.019299
        a02  0.014449  0.013383  0.008850  ...  0.014101  0.017035  0.020008
        a03  0.015605  0.008850  0.005308  ...  0.014903  0.010566  0.01111
        ...

        >>> dfm12
                inner     outer
        a01  0.015710  0.016625
        a02  0.013383  0.014462
        a03  0.005308  0.011291
        ...

        >>> # 使用样本权重计算组间相关性平均值
        >>> dfm13 = x1.rho_default_matrix(aggfunc='sample')[1]
        >>> dfm13
                inner     outer
        a01  0.015710  0.010505
        a02  0.013383  0.008515
        a03  0.005308  0.004487
        ...

        >>> # 其他类型使用方法
        >>> # 利用简单平均权重计算违约相关性矩阵
        >>> dfm21, dfm22 = x1.rho_default_matrix(weight='average')
        >>> # 利用时间相对权重计算违约相关性矩阵
        >>> dfm31, dfm32 = x1.rho_default_matrix(weight='time')
        >>> # 利用时间X配对的样本数作为相对权重计算违约相关性矩阵
        >>> dfm41, dfm42 = x1.rho_default_matrix(weight='time_pair')
        >>> # 利用"时间 X LN(配对的样本数)"作为相对权重计算违约相关性矩阵
        >>> dfm51, dfm52 = x1.rho_default_matrix(weight='time_ln_pair')
        """
        rho_d: np.ndarray = np.zeros(shape=(len(self.__group__), len(self.__group__)))
        for i in range(self.length):
            for j in range(self.length):
                if i > j:
                    rho_d[i, j] = rho_d[j, i]
                elif i == j:
                    rho_d[i, j] = self._inner_rho(i, i, weight, sample_thresh, neg_rho, time_weight)[0]
                else:
                    rho_d[i, j] = self._rho_default_asset(i, j, weight, sample_thresh, neg_rho, time_weight)[1]

        rho_d_avg = np.zeros((self.length, 2))
        rho_d_avg[:, 0] = rho_d.diagonal()
        if aggfunc == 'mean':
            rho_d_avg[:, 1] = (rho_d.sum(axis=1) - rho_d.diagonal()) / (len(rho_d) - 1)
        elif aggfunc == 'sample':
            for i in range(self.length):
                rho_d_avg[i, 1] = self._rho_default_asset(i=i,
                                                          j=None,
                                                          weight=weight,
                                                          thresh=sample_thresh,
                                                          neg_rho=neg_rho,
                                                          time_weight=time_weight)[1]

        dsn_rho = pd.DataFrame(rho_d, columns=self.group, index=self.group)
        dsn_rho_avg = pd.DataFrame(rho_d_avg, columns=['inner', 'outer'], index=self.group)
        return dsn_rho, dsn_rho_avg


class CorrInternalData:
    """
    根据内部违约数据计算违约相关性矩阵、或通过莫顿公式转换为资产相关性矩阵

    Parameters
    ----------
    data：pandas.DataFrame，必须保证有字段“分组”、“日期”、“客户数”、“违约客户数”
        - group：分组名称
        - date: 日期字段
        - customer: 客户数字段
        - default: 违约客户数字段

    >>> from rmwheels.datasets import load_data
    >>> data = load_data("default_history")
    >>> cid = CorrInternalData(data, group='分组', date='年份', customers='客户数', defaults='违约数')
    >>> print(cid.data.head(5))
    """

    def __init__(self, data: pd.DataFrame, **kwargs):
        self.data = data

        # check the input data
        col = ['group', 'date', 'customers', 'defaults']
        data_ = check_data(data, col, **kwargs)

        # initiate the source matrix
        data_pvt = pd.pivot_table(data_, values=['customers', 'defaults'], index='group', columns='date', aggfunc='sum')
        self.customers_array = np.array(data_pvt['customers'])
        self.defaults_array = np.array(data_pvt['defaults'])
        self.group = list(data_pvt.index)
        self.length = len(self.customers_array)

    @staticmethod
    def _weight(arr: np.ndarray):
        return arr / arr.sum()

    def _rho_default_asset(self,
                           i: int,
                           j: int,
                           weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'average'],
                           thresh: int,
                           neg_rho: float | None,
                           time_weight: list[int] | None = None) -> tuple[float, float]:
        """ 计算两组之间的违约相关性 """

        # 根据最小样本数量删除无效的样本
        arr_n1, arr_p1 = self.customers_array[i], self.defaults_array[i]
        if j is not None:
            arr_n2, arr_p2 = self.customers_array[j], self.defaults_array[j]
        else:
            arr_n2 = np.delete(self.customers_array, i, axis=0).sum(axis=0)
            arr_p2 = np.delete(self.defaults_array, i, axis=0).sum(axis=0)

        n_rec = arr_n1.shape[0]
        time_weight.extend([1] * n_rec)
        arr_time = np.array(time_weight)[:n_rec][::-1]

        # 校准由于数据量过小导致的数据缺失，其权重初始值也自动删除
        mask = (arr_n2 >= thresh) * (arr_n1 >= thresh)
        arr_n1, arr_n2 = arr_n1[mask], arr_n2[mask]
        arr_p1, arr_p2 = arr_p1[mask], arr_p2[mask]
        arr_time = arr_time[mask]
        n_rec = arr_n1.shape[0]

        # 确定计量相关性的权重
        if weight == "pair":
            arr_w1 = self._weight(arr_n1 * (arr_n1 - 1))
            arr_w2 = self._weight(arr_n2 * (arr_n2 - 1))
            arr_w3 = self._weight(arr_n1 * arr_n2)
        elif weight == "time":
            arr_w1 = self._weight(arr_time)
            arr_w2, arr_w3 = arr_w1, arr_w1
        elif weight == 'time_pair':
            arr_w1 = self._weight(arr_n1 * (arr_n1 - 1) * arr_time)
            arr_w2 = self._weight(arr_n2 * (arr_n2 - 1) * arr_time)
            arr_w3 = self._weight(arr_n1 * arr_n2 * arr_time)
        elif weight == 'time_ln_pair':
            arr_w1 = self._weight(np.log(arr_n1 * (arr_n1 - 1)) * arr_time)
            arr_w2 = self._weight(np.log(arr_n2 * (arr_n2 - 1)) * arr_time)
            arr_w3 = self._weight(np.log(arr_n1 * arr_n2) * arr_time)
        elif weight == "average":
            arr_w1, arr_w2, arr_w3 = 1 / n_rec, 1 / n_rec, 1 / n_rec
        else:
            raise NotImplementedError(f"weight参数错误，{weight}不存在")

        # 计算相关性
        if i == j:
            arr_pd1 = arr_w1 * (arr_p1 / arr_n1)
            arr_pd1 = arr_pd1.sum()
            arr_jpd = arr_w1 * (arr_p1 / arr_n1) * ((arr_p1 - 1) / (arr_n1 - 1))
            arr_jpd[np.where(arr_p1 == 0)] = 0
            arr_jpd = arr_jpd.sum()
            if arr_pd1 * (1 - arr_pd1) == 0:
                return 0, 0
            else:
                rho = (arr_jpd - arr_pd1 * arr_pd1) / (arr_pd1 * (1 - arr_pd1))
                if np.isnan(rho):
                    return 0, 0
                elif neg_rho is not None and rho < 0:
                    return convert_df_as(arr_pd1, arr_pd1, rho_d=neg_rho)
                else:
                    return convert_df_as(arr_pd1, arr_pd1, rho_d=rho)
        else:
            arr_pd1 = arr_w1 * (arr_p1 / arr_n1)
            arr_pd1 = arr_pd1.sum()
            arr_pd2 = arr_w2 * (arr_p2 / arr_n2)
            arr_pd2 = arr_pd2.sum()
            arr_jpd = arr_w3 * (arr_p1 / arr_n1) * (arr_p2 / arr_n2)
            arr_jpd = arr_jpd.sum()
            if arr_pd1 * (1 - arr_pd1) * arr_pd2 * (1 - arr_pd2) == 0:
                return 0, 0
            else:
                rho = (arr_jpd - arr_pd1 * arr_pd2) / sqrt(arr_pd1 * (1 - arr_pd1) * arr_pd2 * (1 - arr_pd2))
                if np.isnan(rho):
                    return 0, 0
                if neg_rho is not None and rho < 0:
                    return convert_df_as(arr_pd1, arr_pd2, rho_d=neg_rho)

                return convert_df_as(arr_pd1, arr_pd2, rho_d=rho)

    def rho_default_matrix(self,
                           weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'average'] = 'pair',
                           sample_thresh: int = 30,
                           neg_rho: float | None = None,
                           time_weight: list[int | float] | None = None,
                           aggfunc: Literal['mean', 'sample'] = 'mean'):
        """
        给定违约统计数据后计算其对应的违约相关性矩阵以及相关性总体的平均值

        Parameters
        ----------
        weight: str，历史数据权重参数，{'pair','time','time_pair','time_ln_pair','average'}，默认为'pair'
            - pair: 使用样本的配对数作为权重
            - time: 根据时间长短确定权重，默认[3,3,3,2,2,2,1,...]，7年前的权重统一设置为1，若希望改变该权重列表，使用
                    time_weight进行自定义，其顺序是按历史时间由近至远进行分配
            - time_pair: 综合样本配对数与时间长短确定权重
            - time_ln_pair: 综合样本配对数的ln处理与时间长短确定权重
            - average: 每一年的相关性进行简单平均
        sample_thresh: int，计算违约相关性的历史数据最小样本量的设定，默认为30
        neg_rho: float，违约相关性计算结果为负值的时候进行人为修改的取值，默认不修改
        time_weight: list，当权重参数为time、time_pair、time_ln_pair的时候，确定的时间权重列表，默认为[3,3,3,2,2,2,1]
        aggfunc: str，总体相关性平均值的计量方法，默认值为mean
            - mean: 基于违约相关性矩阵，组间相关性为所有与目标分组之间的相关性的平均值
            - sample: 组间相关性是基于将所有非目标分组汇总结果与目标分组，根据相关性计算公式重新计算

        Return
        ------
        dsn_rho: pandas.DataFrame, 包含组名的相关性矩阵数据
        dsn_rho_avg: pandas.DataFrame, 包含组名的相关性平均结果数据

        Examples
        --------
        >>> from rmwheels.datasets import load_data
        >>> data = load_data("default_history")
        >>> x1 = CorrInternalData(data, group='分组', date='年份', cust='客户数', default='违约数')
        >>> # 利用默认设置计算违约相关性、违约相关性平均值
        >>> dfm11, dfm12 = x1.rho_default_matrix()
        >>> dfm11
                  a01       a02       a03  ...       a10       a11       a12
        a01  0.015710  0.014449  0.015605  ...  0.021983  0.016301  0.019299
        a02  0.014449  0.013383  0.008850  ...  0.014101  0.017035  0.020008
        a03  0.015605  0.008850  0.005308  ...  0.014903  0.010566  0.01111
        ...

        >>> dfm12
                inner     outer
        a01  0.015710  0.016625
        a02  0.013383  0.014462
        a03  0.005308  0.011291
        ...

        >>> # 使用样本权重计算组间相关性平均值
        >>> dfm13 = x1.rho_default_matrix(aggfunc='sample')[1]
        >>> dfm13
                inner     outer
        a01  0.015710  0.010505
        a02  0.013383  0.008515
        a03  0.005308  0.004487
        ...

        >>> # 其他类型使用方法
        >>> # 利用简单平均权重计算违约相关性矩阵
        >>> dfm21, dfm22 = x1.rho_default_matrix(weight='average')
        >>> # 利用时间相对权重计算违约相关性矩阵
        >>> dfm31, dfm32 = x1.rho_default_matrix(weight='time')
        >>> # 利用时间X配对的样本数作为相对权重计算违约相关性矩阵
        >>> dfm41, dfm42 = x1.rho_default_matrix(weight='time_pair')
        >>> # 利用"时间 X LN(配对的样本数)"作为相对权重计算违约相关性矩阵
        >>> dfm51, dfm52 = x1.rho_default_matrix(weight='time_ln_pair')
        """
        if time_weight is None:
            time_weight = [3, 3, 3, 2, 2, 2, 1]

        rho_d = np.zeros((self.length, self.length))
        for i in range(self.length):
            for j in range(self.length):
                if i > j:
                    rho_d[i, j] = rho_d[j, i]
                else:
                    rho_d[i, j] = self._rho_default_asset(i, j, weight, sample_thresh, neg_rho, time_weight)[1]

        rho_d_avg = np.zeros((self.length, 2))
        rho_d_avg[:, 0] = rho_d.diagonal()
        if aggfunc == 'mean':
            rho_d_avg[:, 1] = (rho_d.sum(axis=1) - rho_d.diagonal()) / (len(rho_d) - 1)
        elif aggfunc == 'sample':
            for i in range(self.length):
                rho_d_avg[i, 1] = self._rho_default_asset(i=i,
                                                          j=None,
                                                          weight=weight,
                                                          thresh=sample_thresh,
                                                          neg_rho=neg_rho,
                                                          time_weight=time_weight)[1]

        dsn_rho = pd.DataFrame(rho_d, columns=self.group, index=self.group)
        dsn_rho_avg = pd.DataFrame(rho_d_avg, columns=['inner', 'outer'], index=self.group)
        return dsn_rho, dsn_rho_avg

    def rho_asset_matrix(self,
                         weight: Literal['pair', 'time', 'time_pair', 'time_ln_pair', 'average'] = 'pair',
                         sample_thresh: int = 30,
                         neg_rho: float | None = None,
                         time_weight: list[int | float] | None = None,
                         agg: Literal['mean', 'sample'] = 'mean'):
        """
        给定违约统计数据后计算其对应的违约相关性矩阵后，根据莫顿模型推导出资产相关性矩阵

        Parameters
        ----------
        weight: str，历史数据权重参数，{'pair','time','time_pair','time_ln_pair','average'}，默认为'pair'
            - pair: 使用样本的配对数作为权重
            - time: 根据时间长短确定权重，默认[3,3,3,2,2,2,1,...]，7年前的权重统一设置为1，若希望改变该权重列表，使用
                    time_weight进行自定义，其顺序是按历史时间由近至远进行分配
            - time_pair: 综合样本配对数与时间长短确定权重
            - time_ln_pair: 综合样本配对数的ln处理与时间长短确定权重
            - average: 每一年的相关性进行简单平均
        sample_thresh: int，计算违约相关性的历史数据最小样本量的设定，默认为30
        neg_rho: float，违约相关性计算结果为负值的时候进行人为修改的取值，默认不修改
        time_weight: list，当权重参数为time、time_pair、time_ln_pair的时候，确定的时间权重列表，默认为[3,3,3,2,2,2,1]
        agg: str，总体相关性平均值的计量方法，默认值为mean
            - mean: 基于违约相关性矩阵，组间相关性为所有与目标分组之间的相关性的平均值
            - sample: 组间相关性是基于将所有非目标分组汇总结果与目标分组，根据相关性计算公式重新计算

        Return
        ------
        dsn_rho: pandas.DataFrame, 包含组名的相关性矩阵数据
        dsn_rho_avg: pandas.DataFrame, 包含组名的相关性平均结果数据

        Examples
        --------
        >>> dsn = pd.read_excel('examples/test.xlsx', sheet_name='corr_in')
        >>> x1 = CorrInternalData(dsn, group='group', date='date', cust='cust', default='df')
        >>> # 利用默认设置计算资产相关性、资产相关性平均值
        >>> asm11, asm12 = x1.rho_asset_matrix()
        >>> asm11
                  a01       a02       a03  ...       a10       a11       a12
        a01  0.067263  0.062686  0.064558  ...  0.087247  0.069264  0.082764
        a02  0.062686  0.058786  0.038186  ...  0.058487  0.072561  0.086002
        a03  0.064558  0.038186  0.022311  ...  0.059023  0.044644  0.048092
        ...

        >>> asm12
                inner     outer
        a01  0.067263  0.070282
        a02  0.058786  0.062397
        a03  0.022311  0.047499
        ...

        >>> # 使用样本权重计算组间相关性平均值
        >>> asm13 = x1.rho_asset_matrix(agg='sample')[1]
        >>> asm13
                inner     outer
        a01  0.067263  0.044731
        a02  0.058786  0.036838
        a03  0.022311  0.018991
        ...

        >>> # 其他类型使用方法
        >>> # 利用简单平均权重计算资产相关性矩阵
        >>> asm21, asm22 = x1.rho_asset_matrix(weight='average')
        >>> # 利用时间相对权重计算资产相关性矩阵
        >>> asm31, asm32 = x1.rho_asset_matrix(weight='time')
        >>> # 利用时间X配对的样本数作为相对权重计算资产相关性矩阵
        >>> asm41, asm42 = x1.rho_asset_matrix(weight='time_pair')
        >>> # 利用"时间 X LN(配对的样本数)"作为相对权重计算资产相关性矩阵
        >>> asm51, asm52 = x1.rho_asset_matrix(weight='time_ln_pair')
        """
        if time_weight is None:
            time_weight = [3, 3, 3, 2, 2, 2, 1]

        rho_a = np.zeros((self.length, self.length))
        for i in range(self.length):
            for j in range(self.length):
                if i > j:
                    rho_a[i, j] = rho_a[j, i]
                else:
                    rho_a[i, j] = self._rho_default_asset(i=i,
                                                          j=j,
                                                          weight=weight,
                                                          thresh=sample_thresh,
                                                          neg_rho=None,
                                                          time_weight=time_weight)[0]
                    if neg_rho is not None and rho_a[i, j] < 0:
                        rho_a[i, j] = neg_rho

        rho_a_avg = np.zeros((self.length, 2))
        rho_a_avg[:, 0] = rho_a.diagonal()
        if agg == 'mean':
            rho_a_avg[:, 1] = (rho_a.sum(axis=1) - rho_a.diagonal()) / (len(rho_a) - 1)
        elif agg == 'sample':
            for i in range(self.length):
                rho_a_avg[i, 1] = self._rho_default_asset(i=i,
                                                          j=None,
                                                          weight=weight,
                                                          thresh=sample_thresh,
                                                          neg_rho=None,
                                                          time_weight=time_weight)[0]

        dsn_rho = pd.DataFrame(rho_a, columns=self.group, index=self.group)
        dsn_rho_avg = pd.DataFrame(rho_a_avg, columns=['inner', 'outer'], index=self.group)
        return dsn_rho, dsn_rho_avg
